#include "core/common.h"
#include <algorithm>

class ChatServer {
private:
    int serverSocket;
    int port;
    std::vector<std::shared_ptr<ClientInfo>> clients;
    std::mutex clientsMutex;
    std::atomic<bool> running;

public:
    ChatServer(int p = DEFAULT_PORT) : port(p), running(false) {}
    
    ~ChatServer() {
        stop();
    }
    
    bool start() {
        // 创建socket
        serverSocket = socket(AF_INET, SOCK_STREAM, 0);
        if (serverSocket < 0) {
            printLog("ERROR: Failed to create socket");
            return false;
        }
        
        // 设置socket选项
        int opt = 1;
        if (setsockopt(serverSocket, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
            printLog("ERROR: Failed to set socket options");
            return false;
        }
        
        // 绑定地址
        sockaddr_in serverAddr;
        serverAddr.sin_family = AF_INET;
        serverAddr.sin_addr.s_addr = INADDR_ANY;
        serverAddr.sin_port = htons(port);
        
        if (bind(serverSocket, (sockaddr*)&serverAddr, sizeof(serverAddr)) < 0) {
            printLog("ERROR: Failed to bind socket to port " + std::to_string(port));
            return false;
        }
        
        // 开始监听
        if (listen(serverSocket, MAX_CLIENTS) < 0) {    
            printLog("ERROR: Failed to listen on socket");
            return false;
        }
        
        running = true;
        printLog("Chat server started on port " + std::to_string(port));
        printLog("Waiting for connections...");
        
        return true;
    }
    
    void run() {
        while (running) {
            sockaddr_in clientAddr;
            socklen_t clientLen = sizeof(clientAddr);
            
            int clientSocket = accept(serverSocket, (sockaddr*)&clientAddr, &clientLen);
            if (clientSocket < 0) {
                if (running) {
                    printLog("ERROR: Failed to accept client connection");
                }
                continue;
            }
            
            std::string clientIP = inet_ntoa(clientAddr.sin_addr);
            printLog("New connection from " + clientIP);
            
            // 为每个客户端创建处理线程
            std::thread clientThread(&ChatServer::handleClient, this, clientSocket, clientIP);
            clientThread.detach();
        }
    }
    
    void stop() {
        running = false;
        if (serverSocket >= 0) {
            close(serverSocket);
        }
        
        // 关闭所有客户端连接
        std::lock_guard<std::mutex> lock(clientsMutex);
        for (auto client : clients) {
            close(client->socket);
        }
        clients.clear();
        
        printLog("Server stopped");
    }
    
private:
    void handleClient(int clientSocket, const std::string& clientIP) {
        char buffer[BUFFER_SIZE];
        std::string username;
        bool authenticated = false;
        
        while (running) {
            memset(buffer, 0, BUFFER_SIZE);
            int bytesReceived = recv(clientSocket, buffer, BUFFER_SIZE - 1, 0);
            
            if (bytesReceived <= 0) {
                break;
            }
            
            std::string data(buffer);
            Message msg = Message::deserialize(data);
            
            if (!authenticated && msg.type == MSG_LOGIN) {
                username = msg.username;
                if (authenticateUser(username)) {
                    authenticated = true;
                    addClient(clientSocket, username, clientIP);
                    
                    // 发送登录成功消息
                    Message loginSuccess(MSG_LOGIN, "Server", "Login successful");
                    sendToClient(clientSocket, loginSuccess);
                    
                    // 广播用户加入消息
                    Message joinMsg(MSG_CHAT, "Server", username + " joined the chat");
                    broadcastMessage(joinMsg, clientSocket);
                    
                    // 发送在线用户列表
                    sendUserList(clientSocket);
                    
                    printLog("User '" + username + "' logged in from " + clientIP);
                } else {
                    Message loginFailed(MSG_ERROR, "Server", "Username already taken");
                    sendToClient(clientSocket, loginFailed);
                }
            } else if (authenticated) {
                switch (msg.type) {
                    case MSG_CHAT:
                        if (!msg.content.empty()) {
                            Message chatMsg(MSG_CHAT, username, msg.content);
                            broadcastMessage(chatMsg, clientSocket);
                            printLog("Message from " + username + ": " + msg.content);
                        }
                        break;
                        
                    case MSG_USER_LIST:
                        sendUserList(clientSocket);
                        break;
                        
                    case MSG_LOGOUT:
                        printLog("User '" + username + "' logged out");
                        goto cleanup;
                        
                    default:
                        break;
                }
            }
        }
        
cleanup:
        if (authenticated) {
            removeClient(clientSocket);
            Message leaveMsg(MSG_CHAT, "Server", username + " left the chat");
            broadcastMessage(leaveMsg, -1);
        }
        
        close(clientSocket);
        printLog("Client " + clientIP + " disconnected");
    }
    
    bool authenticateUser(const std::string& username) {
        if (username.empty() || username == "Server") {
            return false;
        }
        
        std::lock_guard<std::mutex> lock(clientsMutex);
        for (const auto& client : clients) {
            if (client->username == username) {
                return false;
            }
        }
        return true;
    }
    
    void addClient(int socket, const std::string& username, const std::string& ip) {
        std::lock_guard<std::mutex> lock(clientsMutex);
        clients.push_back(std::make_shared<ClientInfo>(socket, username, ip));
    }
    
    void removeClient(int socket) {
        std::lock_guard<std::mutex> lock(clientsMutex);
        clients.erase(
            std::remove_if(clients.begin(), clients.end(),
                [socket](const std::shared_ptr<ClientInfo>& client) {
                    return client->socket == socket;
                }),
            clients.end()
        );
    }
    
    void sendToClient(int socket, const Message& msg) {
        std::string data = msg.serialize();
        send(socket, data.c_str(), data.length(), 0);
    }
    
    void broadcastMessage(const Message& msg, int excludeSocket = -1) {
        std::string data = msg.serialize();
        
        std::lock_guard<std::mutex> lock(clientsMutex);
        for (const auto& client : clients) {
            if (client->socket != excludeSocket) {
                send(client->socket, data.c_str(), data.length(), 0);
            }
        }
    }
    
    void sendUserList(int socket) {
        std::string userList;
        
        {
            std::lock_guard<std::mutex> lock(clientsMutex);
            for (size_t i = 0; i < clients.size(); ++i) {
                if (i > 0) userList += ", ";
                userList += clients[i]->username;
            }
        }
        
        Message listMsg(MSG_USER_LIST, "Server", "Online users: " + userList);
        sendToClient(socket, listMsg);
    }
};

// 全局服务器实例
ChatServer* server = nullptr;

void signalHandler(int signal) {
    printLog("Received signal " + std::to_string(signal) + ", shutting down...");
    if (server) {
        server->stop();
    }
    exit(0);
}

int main(int argc, char* argv[]) {
    int port = DEFAULT_PORT;
    
    // 解析命令行参数
    if (argc > 1) {
        port = std::atoi(argv[1]);
        if (port <= 0 || port > 65535) {
            std::cerr << "Invalid port number. Using default port " << DEFAULT_PORT << std::endl;
            port = DEFAULT_PORT;
        }
    }
    
    // 设置信号处理
    signal(SIGINT, signalHandler);
    signal(SIGTERM, signalHandler);
    
    // 创建并启动服务器
    server = new ChatServer(port);
    
    if (!server->start()) {
        std::cerr << "Failed to start server" << std::endl;
        delete server;
        return 1;
    }
    
    // 运行服务器
    server->run();
    
    delete server;
    return 0;
} 